from zhipuai import ZhipuAI
import config


def get_format_str(input_string):
    # 删除所有的换行符、制表符和空格
    formatted_string = input_string.replace('\n', '').replace(
        '\r', '').replace('\t', '').replace(' ', '')
    return formatted_string


# 传入用户输入的文本，返回AI回复的文本
def send_to_zhipu(root_text: str) -> str:
    client = ZhipuAI(
        api_key=config.APIKEY)  # 填写自己的APIKey
    response = client.chat.completions.create(
        model="glm-4v-plus",
        stream=False,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": root_text
                    }
                ]
            }
        ]
    )
    return response.choices[0].message.content


def send_with_data(user_text: str, source_data: dict) -> str:
    root_text = f'''
    root_text<
    你现在是一个要进行数据分析的人工智能助手。你需要进行的是抖音直播数据的智能分析。现在我将提供给以下数据集，数据集为字典格式。
    数据集：{source_data}
    请根据数据集中所包含的数据的内容，回复在user_text中的问题。如果user_text中的文本给出了其他数据，也要加以使用。
    请注意，user_text中的文本可能包含多个问题，你需要回答所有的问题。user_text只包含"<<<"">>>"中的内容，不包含root_text中的内容。
    重要提示：请在回答问题时，不要包含root_text中的内容,除了数据集中的数据。
    >
    user_text<<<{user_text}>>>
    '''
    response = send_to_zhipu(get_format_str(root_text))
    return response
